# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/11 10:58 
@Author  : ZhangShenao 
@File    : 2.使用with_fallback实现降级策略.py 
@Desc    : 使用with_fallback()函数,可以实现Runnable的降级策略

例如,当ZhipuAI的LLM调用异常时,自动降级到使用OpenAI的LLM
"""
import os

import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()
prompt = ChatPromptTemplate.from_template('{query}')

# 为LLM组件设置fallback策略
# 当默认的LLM组件调用异常时,自动降级到备选的LLM组件
llm = ChatZhipuAI(model='glm-4-air').with_fallbacks(
    fallbacks=[ChatOpenAI(model='gpt-3.5-turbo', openai_api_base=os.getenv('OPENAI_API_BASE'))],
)
parser = StrOutputParser()

chain = prompt | llm | parser

output = chain.invoke({'query': '你好，你是什么模型？'})
print(output)
